package com.xyf.basicrtc

import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.media.MediaFormat.MIMETYPE_VIDEO_AVC
import android.os.Environment
import android.util.Log
import java.io.FileOutputStream
import java.io.IOException
import java.nio.ByteBuffer

class EncodeH264 constructor(private val width: Int, private val height: Int) {
    lateinit var mediaCodec: MediaCodec
    var onFrame: ((ByteArray) -> Unit)? = null
    var frameIndex = 0
    private var vps_sps_pps_buf: ByteArray? = null

    fun start() {
        try {
            val mediaFormat =
                MediaFormat.createVideoFormat(MIMETYPE_VIDEO_AVC, width, height)
            mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, width * height)
            mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, 15)
            mediaFormat.setInteger(
                MediaFormat.KEY_COLOR_FORMAT,
                MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible
            )
            mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 5)
            mediaCodec = MediaCodec.createEncoderByType(MIMETYPE_VIDEO_AVC)
            mediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
            mediaCodec.start()
        } catch (exc: Exception) {
            Log.e(TAG, "MediaCodec init error", exc)
        }
    }

    fun encodeFrame(nv12: ByteArray) {
        val inputBufferIndex = mediaCodec.dequeueInputBuffer(200000)
        if (inputBufferIndex >= 0) {
            var inputBuffer = mediaCodec.getInputBuffer(inputBufferIndex)
            inputBuffer?.clear()
            inputBuffer?.put(nv12)
            val presentationTimeUs = computePresentationTime(frameIndex.toLong())
            mediaCodec.queueInputBuffer(inputBufferIndex, 0, nv12.size, presentationTimeUs, 0)
            frameIndex++
        }

        val bufferInfo = MediaCodec.BufferInfo()
        var outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, 200000)
        //        推荐 while
        while (outputBufferIndex >= 0) {
            val outputBuffer = mediaCodec.getOutputBuffer(outputBufferIndex)
            outputBuffer?.apply {
                dealFrame(this, bufferInfo)
            }
            mediaCodec.releaseOutputBuffer(outputBufferIndex, false)
            outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, 0)
        }
    }

    fun writeBytes(array: ByteArray?) {
        var writer: FileOutputStream? = null
        try {
            // 打开一个写文件器，构造函数中的第二个参数true表示以追加形式写文件
            writer = FileOutputStream(
                Environment.getExternalStorageDirectory().toString() + "/codec.h265", true
            )
            writer.write(array)
            writer.write('\n'.code)
        } catch (e: IOException) {
            e.printStackTrace()
        } finally {
            try {
                writer?.close()
            } catch (e: IOException) {
                e.printStackTrace()
            }
        }
    }

    private fun computePresentationTime(frameIndex: Long): Long {
        return 240 + frameIndex * 1000000 / 15
    }

    companion object {
        const val TAG = "EncodeH265"
        const val NAL_I = 5
        const val NAL_VPS = 7
    }

    private fun dealFrame(bb: ByteBuffer, bufferInfo: MediaCodec.BufferInfo) {
        var offset = 4
        if (bb[2].toInt() == 0x01) {
            offset = 3
        }
        when ((bb[offset].toInt()) and (0x1f)) {
            NAL_VPS -> {
                vps_sps_pps_buf = ByteArray(bufferInfo.size)
                bb.get(vps_sps_pps_buf)
            }
            NAL_I -> {
                val bytes = ByteArray(bufferInfo.size)
                bb[bytes]
                val newBuf = ByteArray(vps_sps_pps_buf!!.size + bytes.size)
                System.arraycopy(vps_sps_pps_buf, 0, newBuf, 0, vps_sps_pps_buf!!.size)
                System.arraycopy(bytes, 0, newBuf, vps_sps_pps_buf!!.size, bytes.size)
    //            socketLive?.sendData(newBuf)
                onFrame?.run {
                    this(newBuf)
                }
            }
            else -> {
                val bytes = ByteArray(bufferInfo.size)
                bb[bytes]
    //            socketLive?.sendData(bytes)
                onFrame?.run {
                    this(bytes)
                }
            }
        }
    }
}